'''
 * @Author: Benjay·Shaw
 * @Date: 2024-10-31 17:07:50
 * @LastEditors: Benjay·Shaw
 * @LastEditTime: 2024-10-31 22:40:49
 * @Description: 损失函数
'''
import paddle
from networks.vgg import Vgg16


class CACLoss:

    def __init__(self):
        super(CACLoss, self).__init__()
        self.criterion = paddle.nn.CrossEntropyLoss()
        self.criterion_content = paddle.nn.MSELoss()
        self.criterion.cuda(blocking=True)
        self.criterion_content.cuda(blocking=True)
        self.vgg = Vgg16().cuda()

    def __call__(self, img_pre, img_now, result, label):
        ca_map = paddle.nn.functional.sigmoid(x=result[:, 0, :, :] - result
            [:, 1, :, :]).unsqueeze(axis=1).expand_as(y=img_pre)
        img_prev_feature = self.vgg(img_pre * ca_map)
        img_now_feature = self.vgg(img_now * ca_map)
        loss_content = self.criterion_content(img_prev_feature[0],
            img_now_feature[0])
        loss = loss_content + 50 * self.criterion(result, label)
        return loss
